import asyncio
import json
import ssl
from urllib.parse import urlencode
from datetime import datetime

# 引入 mcp 客户端库
from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client

# --- 全局禁用SSL证书验证 ---
try:
    _create_unverified_https_context = ssl._create_unverified_context
except AttributeError:
    pass
else:
    ssl._create_default_https_context = _create_unverified_https_context
# --- ---


async def find_available_tickets_for_train(api_key, train_codes, departure_cities, arrival_cities, query_date=None):
    """
    使用 mcp-server-chinarailway 服务并行查询指定车次在特定日期的余票情况。
    """
    seat_map = {
        '商务座': 'swz_num', '一等座': 'zy_num', '二等座': 'ze_num', '高级软卧': 'gr_num',
        '软卧': 'rw_num', '动卧': 'dw_num', '硬卧': 'yw_num', '软座': 'rz_num',
        '硬座': 'yz_num', '无座': 'wz_num'
    }

    if isinstance(train_codes, (list, tuple, set)):
        wanted_trains = {str(t).strip() for t in train_codes if t}
    else:
        wanted_trains = {str(train_codes).strip()} if train_codes else set()

    if not query_date:
        query_date = datetime.now().strftime('%Y-%m-%d')
    
    print(f"任务开始：并行查询车次 【{', '.join(sorted(wanted_trains))}】 在日期 【{query_date}】 的余票情况。\n")

    base_url = "https://server.smithery.ai/@other-blowsnow/mcp-server-chinarailway/mcp"
    params = {"api_key": api_key}
    url_with_auth = f"{base_url}?{urlencode(params)}"
    
    print(f"将要连接的URL (已包含认证): {base_url}")

    # --- 内部函数：处理单个城市对的查询 ---
    async def query_single_pair(session, dep_city, arr_city):
        """在同一个 session 中查询单个城市对。"""
        print(f"  [任务启动] 查询: {dep_city} -> {arr_city}")
        try:
            result = await session.call_tool(
                "search",
                {
                    "date": query_date,
                    "fromCity": dep_city,
                    "toCity": arr_city
                }
            )
            
            ticket_data = []
            if isinstance(result.content, list) and len(result.content) > 0:
                text_content_obj = result.content[0]
                content_str = text_content_obj.text
                ticket_data = json.loads(content_str)
            
            for train_info in ticket_data:
                actual_train_no = train_info.get("车次")
                if not actual_train_no:
                    continue
                if actual_train_no.strip() in wanted_trains:
                    ticket_list = train_info.get("票列表", {})
                    
                    invalid_vals = {"无", "--", "0", "", "-"}
                    available_seats = {}
                    for seat_name in seat_map.keys():
                        raw = ticket_list.get(seat_name, {}).get("num")
                        if raw is None: continue
                        val = str(raw).strip()
                        if val and val not in invalid_vals:
                            available_seats[seat_name] = val
                    
                    start_time, end_time = None, None
                    start_keys = ['出发时间', '始发时间', '发车时间', '出发时刻', 'start_time', 'departure_time', '发车']
                    end_keys = ['到达时间', '终到时间', '到站时间', '到达时刻', 'arrival_time', 'end_time', '到达']
                    for k in start_keys:
                        if v := train_info.get(k): start_time = str(v).strip(); break
                    for k in end_keys:
                        if v := train_info.get(k): end_time = str(v).strip(); break

                    if available_seats:
                        ticket_info = {
                            "日期": query_date, "车次": actual_train_no, "出发地": dep_city,
                            "到达地": arr_city, "出发时间": start_time or "", "到达时间": end_time or "",
                            "余票席位": available_seats
                        }
                        print(f"  --> ✅ [查询成功] {dep_city} -> {arr_city} | 车次 {actual_train_no} 发现余票: {available_seats}")
                        return ticket_info
                    else:
                        print(f"  --> 🟡 [已找到但无票] {dep_city} -> {arr_city} | 车次 {actual_train_no} 当前无余票。")
                        return None # 明确返回 None
            
            # print(f"  --> ❌ [未找到车次] 在 {dep_city} -> {arr_city} 的线路中未找到任何目标车次。")
            return None # 明确返回 None

        except json.JSONDecodeError as e:
            print(f"  --> 🚨 [JSON解析错误] {dep_city} -> {arr_city}: {e}. 原始返回内容: {result.content}")
        except Exception as e:
            print(f"  --> 🚨 [查询异常] {dep_city} -> {arr_city}: {e}")
        return None # 出现任何异常都返回 None

    # --- 主逻辑：建立连接，创建并执行并行任务 ---
    all_results = []
    try:
        async with streamablehttp_client(url_with_auth) as (read, write, _):
            async with ClientSession(read, write) as session:
                print("正在初始化与 MCP 服务器的连接...")
                await session.initialize()
                print("✅ 服务器初始化成功！\n")
                
                try:
                    tools_result = await session.list_tools()
                    tool_names = [t.name for t in tools_result.tools]
                    print(f"✅ 服务器可用工具列表: {tool_names}\n")
                    if "search" not in tool_names:
                         print(f"❌ 警告：服务器工具列表中未找到 'search' 工具。程序将退出。")
                         return []
                except Exception as e:
                    print(f"❌ 获取工具列表失败: {e}")
                    return []

                # 创建所有查询任务
                tasks = []
                for dep_city in departure_cities:
                    for arr_city in arrival_cities:
                        if dep_city == arr_city:
                            continue
                        # 将对内部函数的调用作为一个任务添加到列表
                        tasks.append(query_single_pair(session, dep_city, arr_city))
                
                # 并发执行所有任务
                print(f"准备并发执行 {len(tasks)} 个查询任务...")
                results_from_gather = await asyncio.gather(*tasks, return_exceptions=True)
                
                # 处理并收集结果
                for res in results_from_gather:
                    if isinstance(res, Exception):
                        # 如果任务返回的是一个异常，打印它
                        print(f"  --> 🚨 [任务执行异常]: {res}")
                    elif res:
                        # 如果任务返回了有效的车票信息，添加到最终列表
                        all_results.append(res)

    # === 【重要修改】 增强错误捕获逻辑 ===
    except ExceptionGroup as eg:
        print(f"❌ 连接或初始化 MCP 会话时发生并发错误。包含 {len(eg.exceptions)} 个子异常:")
        for i, e in enumerate(eg.exceptions):
            print(f"  - 子异常 {i+1}: {type(e).__name__} - {e}")
    # =======================================
    except asyncio.exceptions.CancelledError:
        print("任务被用户取消。")
    except Exception as e:
        # 保留通用的异常捕获，以防万一
        print(f"❌ 连接或初始化 MCP 会话失败: {type(e).__name__} - {e}")
        # 这里可以保留之前的提示，因为URL错误等仍然可能发生
        return []

    return all_results

# --- 使用示例 ---
async def main():
    # ==================== 请在这里修改您要查询的信息 ====================
    my_api_key = "cff0ac8f-51ef-4b46-ae98-3b1875b59ddb"
    target_date = "2025-10-01"
    target_train = ["G2691","G2694"]
    departure_points = ["潮汕", "普宁", "陆丰", "惠东南", "深圳北"]
    arrival_points = [  "太湖南", "安庆","郑州东"]
    # =================================================================

    if not my_api_key or "YOUR_API_KEY" in my_api_key:
        print("❌ 错误：请在 main 函数中设置您的 my_api_key。")
        return

    target_trains = [t.strip() for t in (target_train if isinstance(target_train, (list, tuple, set)) else [target_train]) if t]

    results = await find_available_tickets_for_train(
        my_api_key, target_trains, departure_points, arrival_points, target_date
    )

    print("\n---------- ✅ 查询任务完成，结果汇总如下 ----------")
    if results:
        # 为了美观，可以按出发地和车次排序
        sorted_results = sorted(results, key=lambda x: (x['出发地'], x['车次']))
        print(json.dumps(sorted_results, ensure_ascii=False, indent=4))
    else:
        print("所有查询组合均未发现符合条件的余票。")

if __name__ == '__main__':
    asyncio.run(main())